Benchmark: meta-init + rank0-broadcast loading vs main (8x H200, cold page cache)

Measured on the real FSDPTrainRayActor.init model-loading path;
checkpoint page cache evicted (posix_fadvise DONTNEED) before every run;
latency and peak memory are the max across ranks.
Per-rank shard sha256 digests are bitwise-identical between the two
modes for all four families (same world size, same sharding plan, so the
local shards must match exactly).

Init latency (s) — lower is better

xychart-beta
    title "init latency (s): main (back) vs stream (front)"
    x-axis ["SD3.5-medium", "Qwen-Image", "Wan2.2-A14B (2 experts)", "LTX-2.3"]
    y-axis "seconds" 0 --> 434
    bar [37.2, 247.6, 377.0, 68.2]
    bar [15.8, 120.2, 133.2, 56.3]

Peak GPU memory per rank (GB) — lower is better

xychart-beta
    title "peak GPU/rank (GB): main (back) vs stream (front)"
    x-axis ["SD3.5-medium", "Qwen-Image", "Wan2.2-A14B (2 experts)", "LTX-2.3"]
    y-axis "GB" 0 --> 88
    bar [9.3, 76.3, 60.1, 35.5]
    bar [2.9, 10.9, 14.6, 5.2]

xychart-beta
    title "init speedup (x, higher is better)"
    x-axis ["SD3.5-medium", "Qwen-Image", "Wan2.2-A14B (2 experts)", "LTX-2.3"]
    y-axis "speedup (x)" 0 --> 3
    bar [2.4, 2.1, 2.8, 1.2]

xychart-beta
    title "peak GPU/rank reduction (%)"
    x-axis ["SD3.5-medium", "Qwen-Image", "Wan2.2-A14B (2 experts)", "LTX-2.3"]
    y-axis "% reduction" 0 --> 99
    bar [68.8, 85.7, 75.7, 85.4]

LoRA load equivalence (same harness, --lora, rank 64, gaussian init)

Both branches wrap each component with PEFT before FSDP: main inits the
adapter on every rank (identical seed), this PR inits it on rank0 only and
broadcasts A/B together with the base weights (non-rank0 builds the adapter
on meta via get_peft_model(..., low_cpu_mem_usage=True)). Checks per run:
no meta residue, lora_B == 0, lora_A finite and nonzero, buffers
replicated across ranks; then per-rank digests are compared across branches.

Base weights are bitwise-identical to main under the PEFT-wrapped FQNs for
all four families. lora_A even matches bitwise wherever both branches
consume the RNG in the same order; where it differs it is still the same
N(0, 1/r) init from the same peft code, and lora_B == 0 makes the
initial model function identical to the base model either way.

Seed determinism: the whole LoRA suite was executed twice with the same
seed; run1 and run2 full digests (base + random-init lora_A) are
bitwise-identical on every rank, for every family, on both branches.

Loading mechanics: low_cpu_mem_usage, meta init, and _keep_in_fp32_modules

This is the most confusing corner of the PR, so spelling it out.

diffusers' from_pretrained has two loading paths:

• low_cpu_mem_usage=True (diffusers' default): self-contained fast path.
  diffusers builds the module on its own internal meta device, then reads the
  checkpoint and materializes every param directly by assignment. It always
  produces real weights and bypasses any ambient init context (assignment does
  not go through parameter registration, so an external init_empty_weights
  cannot intercept it). Peak CPU ~= 1x model size.
• low_cpu_mem_usage=False: plain path. Runs the model's ordinary
  __init__ (params allocated through normal registration), then reads the
  full state dict and copy_()s it into the params. In normal use this peaks
  at ~2x model size -- that is exactly what True was invented to avoid.

Why each rank uses what it uses:

Non-rank0 must use False: it is the only path whose construction step can
be intercepted by our ambient meta context. With True, diffusers would hand
back fully materialized weights on every rank and the meta-shell design (shard
on meta -> to_empty only this rank's shard on GPU -> broadcast from rank0)
would silently degenerate into N full copies; the actor guards this with an
explicit "did not honor meta initialization" check.

The conflict: some model classes (e.g. WanTransformer3DModel) declare
_keep_in_fp32_modules -- submodules that must stay fp32 even when the rest
loads in bf16. diffusers implements that pinning only inside the True
path, and rather than silently dropping the pin it raises up-front:
ValueError: low_cpu_mem_usage cannot be False when keep_in_fp32_modules is True. So non-rank0 (forced False) cannot load such a class as-is.

How we handle it: on non-rank0 only, _keep_in_fp32_modules is temporarily
cleared on the class for the duration of the from_pretrained call. This is
safe because non-rank0 params are meta placeholders anyway -- authoritative
dtypes come from rank0 (which loads with True and honors the pin), and
sync_model_dtypes broadcasts rank0's per-tensor dtypes to every rank before
sharding, so all shards agree.

The Wan crash this benchmark surfaced (fixed in this PR)

Components used to be loaded through DiffusionPipeline.from_pretrained with
every non-target component passed as None. That None-skip convention only
works for components that are required in the pipeline __init__ signature:
only those land in passed_class_obj, the map whose None entries suppress
loading. WanPipeline declares transformer/transformer_2 as optional
with a None default, so _get_signature_keys drops them from
expected_modules, the None we passed was silently discarded, and both 14B
experts loaded from disk on every rank -- during load_scheduler too. On
non-rank0 that load ran with low_cpu_mem_usage=False while the fp32-guard
patch had been applied to the scheduler's class, so the unpatched
WanTransformer3DModel hit the ValueError above and every multi-rank init
crashed. Single-GPU runs never crashed (rank0 loads with True); the waste
was silent.

Fix: resolve the component class from model_index.json and load it directly
via cls.from_pretrained(ckpt, subfolder=component). Only the named component
is ever touched, so there is nothing to skip, and the fp32 patch provably
targets the class actually being loaded.

🤖 Generated with Claude Code